import time
import threading
import zmq
import random


# ======================
# 请求-响应模式 (REQ-REP)
# ======================
def request_response_server():
    """请求-响应模式 - 服务器端"""
    context = zmq.Context()
    socket = context.socket(zmq.REP)
    socket.bind("tcp://*:5555")

    print("REP服务器已启动，等待请求...")
    while True:
        # 接收请求
        message = socket.recv_string()
        print(f"收到请求: {message}")

        # 处理请求
        time.sleep(0.5)  # 模拟处理时间
        response = f"已处理: {message}"

        # 发送响应
        socket.send_string(response)


def request_response_client():
    """请求-响应模式 - 客户端"""
    context = zmq.Context()
    socket = context.socket(zmq.REQ)
    socket.connect("tcp://localhost:5555")

    for i in range(3):
        # 发送请求
        message = f"请求 #{i + 1}"
        socket.send_string(message)
        print(f"客户端发送: {message}")

        # 接收响应
        response = socket.recv_string()
        print(f"收到响应: {response}")
        time.sleep(1)


# ======================
# 发布-订阅模式 (PUB-SUB)
# ======================
def publisher_server():
    """发布-订阅模式 - 发布者"""
    context = zmq.Context()
    socket = context.socket(zmq.PUB)
    socket.bind("tcp://*:5556")

    topics = ["news", "sports", "weather"]
    print("PUB服务器已启动，开始发布消息...")

    while True:
        # 随机选择一个主题和消息
        topic = random.choice(topics)
        message = f"{topic}更新: {time.strftime('%H:%M:%S')}"

        # 发布消息 (格式: "主题 消息内容")
        socket.send_string(f"{topic} {message}")
        print(f"发布消息: {topic} - {message}")
        time.sleep(1)


def subscriber_client(topic):
    """发布-订阅模式 - 订阅者"""
    context = zmq.Context()
    socket = context.socket(zmq.SUB)
    socket.connect("tcp://localhost:5556")

    # 订阅特定主题
    socket.setsockopt_string(zmq.SUBSCRIBE, topic)
    print(f"订阅者已启动，订阅主题: {topic}")

    while True:
        # 接收消息
        message = socket.recv_string()
        print(f"订阅者收到消息: {message}")
        time.sleep(0.5)


# ======================
# 推送-拉取模式 (PUSH-PULL)
# ======================
def push_server():
    """推送-拉取模式 - 推送者"""
    context = zmq.Context()
    socket = context.socket(zmq.PUSH)
    socket.bind("tcp://*:5557")

    print("PUSH服务器已启动，开始推送任务...")

    # 发送开始信号
    print("发送开始信号...")
    socket.send(b"START")

    for i in range(5):
        # 推送任务
        task = f"任务 #{i + 1}"
        socket.send_string(task)
        print(f"推送任务: {task}")
        time.sleep(0.5)


def pull_worker(worker_id):
    """推送-拉取模式 - 工作者"""
    context = zmq.Context()
    socket = context.socket(zmq.PULL)
    socket.connect("tcp://localhost:5557")

    print(f"工作者 #{worker_id} 已启动，等待任务...")

    # 等待开始信号
    signal = socket.recv()
    print(f"工作者 #{worker_id} 收到信号: {signal.decode()}")

    while True:
        # 拉取任务
        task = socket.recv_string()
        print(f"工作者 #{worker_id} 收到任务: {task}")

        # 处理任务
        time.sleep(random.uniform(0.5, 2.0))
        print(f"工作者 #{worker_id} 完成任务: {task}")


# ======================
# 路由-代理模式 (ROUTER-DEALER)
# ======================
def router_server():
    """路由-代理模式 - 路由器"""
    context = zmq.Context()
    socket = context.socket(zmq.ROUTER)
    socket.bind("tcp://*:5558")

    print("ROUTER服务器已启动，等待连接...")

    workers = {}  # 存储工作者ID和最后活动时间

    while True:
        # 接收消息 (包含身份帧和消息帧)
        identity, _, message = socket.recv_multipart()
        worker_id = identity.decode()

        if message == b"READY":
            # 新工作者注册
            print(f"工作者 {worker_id} 已注册")
            workers[worker_id] = time.time()
        else:
            # 处理工作结果
            print(f"收到来自 {worker_id} 的结果: {message.decode()}")

            # 分配新任务
            task = f"新任务-{time.strftime('%H%M%S')}"
            socket.send_multipart([identity, b"", task.encode()])
            print(f"向 {worker_id} 分配任务: {task}")


def dealer_worker(worker_id):
    """路由-代理模式 - 工作者"""
    context = zmq.Context()
    socket = context.socket(zmq.DEALER)
    socket.setsockopt_string(zmq.IDENTITY, worker_id)
    socket.connect("tcp://localhost:5558")

    print(f"工作者 {worker_id} 已启动")

    # 发送就绪信号
    socket.send_multipart([b"", b"READY"])

    while True:
        # 接收任务
        _, message = socket.recv_multipart()
        task = message.decode()
        print(f"工作者 {worker_id} 收到任务: {task}")

        # 处理任务
        time.sleep(random.uniform(0.5, 1.5))
        result = f"{task} 完成于 {time.strftime('%H:%M:%S')}"

        # 发送结果
        socket.send_multipart([b"", result.encode()])
        print(f"工作者 {worker_id} 发送结果: {result}")


# ======================
# 主程序
# ======================
if __name__ == "__main__":
    # 启动所有服务器
    # threading.Thread(target=request_response_server, daemon=True).start()
    # threading.Thread(target=publisher_server, daemon=True).start()
    # threading.Thread(target=push_server, daemon=True).start()
    threading.Thread(target=router_server, daemon=True).start()

    # 等待服务器启动
    # time.sleep(1)

    # print("\n===== 测试请求-响应模式 =====")
    # request_response_client()
    #
    # print("\n===== 测试发布-订阅模式 =====")
    # 启动多个订阅者
    # threading.Thread(target=subscriber_client, args=("news",), daemon=True).start()
    # threading.Thread(target=subscriber_client, args=("sports",), daemon=True).start()
    # threading.Thread(target=subscriber_client, args=("weather",), daemon=True).start()
    #
    # # 给订阅者时间连接
    # time.sleep(2)
    #
    # print("\n===== 测试推送-拉取模式 =====")
    # 启动多个工作者
    # threading.Thread(target=pull_worker, args=("1",), daemon=True).start()
    # threading.Thread(target=pull_worker, args=("2",), daemon=True).start()
    #
    # 给工作者时间连接
    time.sleep(2)
    #
    print("\n===== 测试路由-代理模式 =====")
    # 启动多个工作者
    threading.Thread(target=dealer_worker, args=("Worker-A",), daemon=True).start()
    threading.Thread(target=dealer_worker, args=("Worker-B",), daemon=True).start()

    # 保持主线程运行
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\n程序已终止")
